#include "calibration.h"
#include <vector>
#include <QDebug>

using namespace std;
using namespace cv;

Calibration::Calibration(QObject* parent)
	: captureImages(nullptr)
{
}

Calibration::~Calibration()
{
}

void Calibration::initBoard(const QString& fileName)
{
	FileStorage fs(fileName.toStdString(), FileStorage::READ);
	int width, height;
	double squareSize;
	fs["board_width"] >> width;
	fs["board_height"] >> height;
	fs["square_size"] >> squareSize;
	fs.release();

	this->boardSize = {width, height};
	this->squareSize = squareSize;
}

void Calibration::setImageList(QList<cv::Mat>* list)
{
	captureImages = list;
}

void Calibration::save(const QString& fileName)
{
	FileStorage fs(fileName.toStdString(), FileStorage::WRITE);
	fs << "camera_matrix" << cameraMatrix;
	fs << "dist_coeffs" << distCoeffs;
	fs.release();
}

void Calibration::load(const QString& fileName)
{
	FileStorage fs(fileName.toStdString(), FileStorage::READ);
	fs["camera_matrix"] >> cameraMatrix;
	fs["dist_coeffs"] >> distCoeffs;
	fs.release();
}


void Calibration::run()
{
	if (captureImages->size() < 10) {
		emit processed(false, "capture image number is too less");
		return;
	}

	// 寻找标定板关键点
	vector<vector<Point2f>> imagePoints;
	findCorners(imagePoints);
	if (imagePoints.empty()) {
		emit processed(false, "find chess board corners failed");
		return;
	}
	qDebug() << "calibration finished";

	// 标定
	Size imageSize = captureImages->begin()->size();
	if (!runCalibration(imageSize, imagePoints)) {
		emit processed(false, "caliration image failed");
		return;
	}

	// 矫正原图像
	for (auto& captureImage : *captureImages) {
		Mat correctImage;
		correct(captureImage, correctImage);
		imshow(srcImageTitle, captureImage);
		imshow(winTitle, correctImage);
		waitKey(100);
	}
	destroyWindow(srcImageTitle);
	destroyWindow(winTitle);

	emit processed(true, "");
}

void Calibration::syncRun()
{
	run();
}


void Calibration::findCorners(std::vector<std::vector<cv::Point2f>>& imagePoints)
{
	for (int i = 0; i < captureImages->size(); ++i) {
		Mat view;
		captureImages->at(i).copyTo(view);
		vector<Point2f> pointBuf;
		bool found = findChessboardCorners(view, boardSize, pointBuf,
			CALIB_CB_ADAPTIVE_THRESH | CALIB_CB_FAST_CHECK | CALIB_CB_NORMALIZE_IMAGE);

		if (found) {
			Mat viewGray;
			if (view.type() != CV_8UC1) {
				cvtColor(view, viewGray, COLOR_BGR2GRAY);
			} else {
				viewGray = view;
				cvtColor(view, view, COLOR_GRAY2BGR);
			}
			cornerSubPix(viewGray, pointBuf, Size(11, 11),
				Size(-1, -1), TermCriteria(TermCriteria::EPS + TermCriteria::MAX_ITER, 30, 0.1));
			imagePoints.push_back(pointBuf);

			drawChessboardCorners(view, boardSize, Mat(pointBuf), found);
		}

		string msg = format("%d/%d %s", i + 1, captureImages->size(), found ? "found" : "no found");
		int baseLine;
		Size textSize = getTextSize(msg, 1, 1, 1, &baseLine);
		Point textOrigin(view.cols - 2 * textSize.width - 10, view.rows - 2 * baseLine - 10);

		putText(view, msg, textOrigin, 1, 1, found ? Scalar(0, 255, 0) : Scalar(0, 0, 255));

		imshow(winTitle, view);
		waitKey(10);
	}
}

bool Calibration::runCalibration(Size& imageSize, std::vector<std::vector<cv::Point2f>>& imagePoints)
{
	vector<Mat> rvecs, tvecs;
	vector<float> reprojErrs;

	cameraMatrix = Mat::eye(3, 3, CV_64F);
	cameraMatrix.at<double>(0, 0) = 1.0;

	distCoeffs = Mat::zeros(8, 1, CV_64F);

	vector<vector<Point3f> > objectPoints(1);
	for (int i = 0; i < boardSize.height; ++i) {
		for (int j = 0; j < boardSize.width; ++j) {
			objectPoints[0].push_back(Point3f(float(j * squareSize), float(i * squareSize), 0));
		}
	}

	objectPoints.resize(imagePoints.size(), objectPoints[0]);

	int flags = CALIB_FIX_PRINCIPAL_POINT | CALIB_ZERO_TANGENT_DIST | 
		CALIB_FIX_ASPECT_RATIO | CALIB_FIX_K4 | CALIB_FIX_K5;
	double rms = calibrateCamera(objectPoints, imagePoints, imageSize, cameraMatrix,
		distCoeffs, rvecs, tvecs, flags);

	qDebug() << "Re-projection error reported by calibrateCamera: " << rms;

	bool ok = checkRange(cameraMatrix) && checkRange(distCoeffs);

	vector<Point2f> imagePoint;
	int totalPoints = 0;
	double totalErr = 0, err;
	reprojErrs.resize(objectPoints.size());

	for (int i = 0; i < (int)objectPoints.size(); ++i)
	{
		projectPoints(Mat(objectPoints[i]), rvecs[i], tvecs[i], cameraMatrix,
			distCoeffs, imagePoint);
		err = norm(Mat(imagePoints[i]), Mat(imagePoint), NORM_L2);

		int n = (int)objectPoints[i].size();
		reprojErrs[i] = (float)std::sqrt(err * err / n);
		totalErr += err * err;
		totalPoints += n;
	}

	double totalAvgErr = std::sqrt(totalErr / totalPoints);
	qDebug() << "Total average error reported by calibrateCamera: " << totalAvgErr;

	return ok;
}

void Calibration::correct(cv::Mat& inMat, cv::Mat& outMat)
{
	if (inMat.empty()) {
		return;
	}
	Mat map1, map2;
	Size imageSize = inMat.size();
	initUndistortRectifyMap(cameraMatrix, distCoeffs, Mat(),
		getOptimalNewCameraMatrix(cameraMatrix, distCoeffs, imageSize, 1, imageSize, 0),
		imageSize, CV_16SC2, map1, map2);

	remap(inMat, outMat, map1, map2, INTER_LINEAR);
}

